
'use client'
import { AutoComplete, Button, Card, Form, Image, Input, InputNumber, InputTag, Message, Radio, Select, Space, Switch } from '@arco-design/web-react'
import React, { useEffect, useState } from 'react'
import PromptInput from '../../../components/PromptInput'
import { models } from '../../../constants/models'
import Images from '../../../components/Images'
import { useLocalStorageState } from 'ahooks'
import { useTRPCClient } from '../../server/trpc/client'
import { useStorageForm } from '../../../hooks/useStorageForm'
const RadioGroup = Radio.Group;

const FormItem = Form.Item
const Option = Select.Option

type FormData = {
  model: string,
  promptText: string,
  shape: string,
  isHyper: boolean,
}

const FORM_DATA_STORAGE_KEY = 'history_form'
const HISTORY_IMAGE_STORAGE_KEY = 'history_images'

const IMAGE_SHAPES = [
  {
    label: '方形',
    value: 'square',
    size: {
      height: 1024,
      width: 1024,
    }
  },
  {
    label: '横屏',
    value: 'landscape',
    size: {
      height: 832,
      width: 1216,
    }
  },
  {
    label: '竖屏',
    value: 'portrait',
    size: {
      height: 1216,
      width: 832,
    }
  }
]

const Page = () => {
  // const [form] = Form.useForm<FormData>()
  const { form, saveForm } = useStorageForm<FormData>({
    storageKey: FORM_DATA_STORAGE_KEY,
    initValue: {
      model: models[0]?.model,
      promptText: '',
      shape: 'square',
      isHyper: true,
    }
  })
  const [historyImages, setHistoryImages] = useLocalStorageState<string[]>(HISTORY_IMAGE_STORAGE_KEY, { defaultValue: [] })
  const [loading, setLoading] = useState(false)
  const { image } = useTRPCClient();

  const handleSubmit = async (params: FormData) => {
    saveForm()

    setLoading(true);
    const queueCount = await image.getQueueCount.query();
    if (queueCount) {
      Message.info(`排队中，预计剩余时间${queueCount * 15}秒`)
    }

    try {
      const selectedShape = IMAGE_SHAPES.find(shape => shape.value === params.shape);
      const apiParams = {
        ...params,
        promptText: params.promptText,
        width: selectedShape?.size.width || 1024,
        height: selectedShape?.size.height || 1024
      };
      const res = await image.text2img.mutate(apiParams)
      console.log('success');
      if (res) {
        const newImages = res.concat(historyImages)
        setHistoryImages(newImages);
      }
    } catch (e) {
      console.log(e);
      
      Message.error('生成失败，请重试');
    } finally {

      setLoading(false);

    }
  }

  return (
    <Card title="文生图" className={"max-w-200"}>
      <Form layout="vertical" form={form} onSubmit={handleSubmit}>
        <FormItem rules={[{ required: true }]} required label='模型' field="model">
          <Select
            placeholder='请选择模型'
          >
            {models.map((option) => (
              <Option key={option.model} value={option.model}>
                {option.name}
              </Option>
            ))}
          </Select>
        </FormItem>


        <FormItem rules={[{ required: true }]} required label='提示词' field="promptText">
          <PromptInput />
        </FormItem>

        <FormItem rules={[{ required: true }]} required label='图片形状' field="shape">
          <RadioGroup
            size="large"
            type='button'
          >
            {IMAGE_SHAPES.map((shape) => (
              <Radio key={shape.value} value={shape.value}>
                {shape.label}
              </Radio>
            ))}
          </RadioGroup>
        </FormItem>
        <FormItem rules={[{ required: true }]} triggerPropName='checked' label='加速生成' extra="加速生成的图片质量略微有下降" field="isHyper">
          <Switch />
        </FormItem>
        <FormItem>
          <Button shape="round" long loading={loading} size="large" type='primary' htmlType='submit'>生成</Button>
        </FormItem>
      </Form>
      <Images data={historyImages} />
    </Card>
  )
}

export default Page
